# 划分数据集
import os
import shutil
import random

# 原始数据集路径
SRC_DIR = 'dataset'
# 输出路径
DST_DIR = 'data1'
TRAIN_DIR = os.path.join(DST_DIR, 'train')
VAL_DIR = os.path.join(DST_DIR, 'val')
SPLIT_RATIO = 0.7  # 训练集比例

random.seed(42)  # 保证可复现

# 创建目标文件夹
for split in ['train', 'val']:
    split_dir = os.path.join(DST_DIR, split)
    if not os.path.exists(split_dir):
        os.makedirs(split_dir)

# 遍历每个类别
for class_name in os.listdir(SRC_DIR):
    class_path = os.path.join(SRC_DIR, class_name)
    if not os.path.isdir(class_path):
        continue
    images = os.listdir(class_path)
    random.shuffle(images)
    n_train = int(len(images) * SPLIT_RATIO)
    train_images = images[:n_train]
    val_images = images[n_train:]

    # 创建类别子文件夹
    train_class_dir = os.path.join(TRAIN_DIR, class_name)
    val_class_dir = os.path.join(VAL_DIR, class_name)
    os.makedirs(train_class_dir, exist_ok=True)
    os.makedirs(val_class_dir, exist_ok=True)

    # 复制训练集图片
    for img in train_images:
        src_img = os.path.join(class_path, img)
        dst_img = os.path.join(train_class_dir, img)
        shutil.copy2(src_img, dst_img)

    # 复制验证集图片
    for img in val_images:
        src_img = os.path.join(class_path, img)
        dst_img = os.path.join(val_class_dir, img)
        shutil.copy2(src_img, dst_img)

    print(f'类别 {class_name} 处理完成，训练集 {len(train_images)} 张，验证集 {len(val_images)} 张')

print('数据集划分完成，输出目录为 data1/train 和 data1/val')